# -*- coding: utf-8 -*-

"""
哈夫曼编码、解码，实现了统计文本中字符频率，并根据频率信息生成哈夫曼树，
并且提供了编码，解码的功能。
"""
import heapq

class HuffmanTree(object):
    """
    哈夫曼树
    """
    def __init__(self, string):
        """
        初始化哈夫曼数据，根据文本生成哈夫曼树，并生成哈夫曼编码表。
        @param string: 要进行编码的文本
        """
        self.original_text = string
        self.cipher_text = ''
        self.char_frequency_dict = {}
        self.tree = []
        self.encode_table = {}
        
        if len(string) is not 0:
            self._cal_char_frequency(string)
            self._make_huffman_tree()
            self._generate_huffman_encode_table(self.tree)
    
    def _cal_char_frequency(self, string):
        """统计各个字符出现的次数（按照8bit一个字符转换为0-255的数值）"""
        for char in string:
            code = ord(char)    # 获得Ascii码
            if code in self.char_frequency_dict:
                self.char_frequency_dict[code] += 1
            else:
                self.char_frequency_dict[code] = 1
        
    def _make_huffman_tree(self):
        """生成哈夫曼树"""
        # 一个节点有两部分组成，第一部分是频率，第二部分是Ascii码     
        self.tree = [[self.char_frequency_dict[char], char] for char in self.char_frequency_dict]
        heapq.heapify(self.tree)
        while len(self.tree) > 1:
            childR, childL = heapq.heappop(self.tree), heapq.heappop(self.tree)
            parent = [childL[0] + childR[0], childL, childR]
            heapq.heappush(self.tree, parent)
        
        #去掉list，转换成元组
        self.tree = self.tree[0]
        
    def _generate_huffman_encode_table(self, huffman_tree, prefix=''):
        """生成哈夫曼编码表"""
        if len(huffman_tree) is 2:
            self.encode_table[huffman_tree[1]] = prefix
        else:
            self._generate_huffman_encode_table(huffman_tree[1], prefix + '0')
            self._generate_huffman_encode_table(huffman_tree[2], prefix + '1')

    def encode(self):
        """加密，并返回加密后的文本
        @return: 加密后的文本
        """
        encode_string = ""
        for char in self.original_text:
            code = ord(char)
            encode_string += self.encode_table[code]
        
        self.cipher_text = encode_string
        return encode_string
    
    def decode(self):
        """解密，并返回解密后的文本
        @return: 解密后的文本
        """
        decode_string = ""
        trees = self.tree
        if trees is None:
            return
        
        for bit in self.cipher_text:
            if bit is '0':
                childL = trees[1]
                trees = childL
                if len(childL) == 2:
                    decode_string += chr(childL[1])    #这种方法，暂时不支持中文
                    trees = self.tree
            else:
                childR = trees[2]
                trees = childR
                if len(childR) == 2:
                    decode_string += chr(childR[1])    #这种方法，暂时不支持中文
                    trees = self.tree
        return decode_string

if __name__ == "__main__":
    string = "Hello World!"
    huffman_tree = HuffmanTree(string)
    print huffman_tree.char_frequency_dict
    print huffman_tree.tree
    print huffman_tree.encode_table
    print huffman_tree.encode()
    print huffman_tree.decode()    

